﻿#pragma once

#include <learn/geom/vec.h>

#include <functional>

namespace xi
{

namespace geom
{

template <typename T> struct Quat
{
    using Vec = VecT<3, T>;

    T s;
    Vec v;

    explicit Quat() = default;

    /**
     * @brief 构造四元数
     *
     * @param[in] s
     * @param[in] v
     */
    Quat(T s, const Vec &v) : s(s), v(v)
    {
    }

    /**
     * @brief 构造四元数
     *
     * @param[in] a
     * @param[in] b
     * @param[in] c
     * @param[in] d
     */
    Quat(T a, T b, T c, T d) : s(a), v(b, c, d)
    {
    }

#define QUAT_BINARY_OP(op)                                                                                             \
    Quat operator op(const Quat &q) const                                                                              \
    {                                                                                                                  \
        return Quat(s op q.s, v op q.v);                                                                               \
    }                                                                                                                  \
    const Quat &operator op##=(const Quat &q)                                                                          \
    {                                                                                                                  \
        s op## = q.s;                                                                                                  \
        v op## = q.v;                                                                                                  \
        return *this;                                                                                                  \
    }

    QUAT_BINARY_OP(+)
    QUAT_BINARY_OP(-)
#undef QUAT_BINARY_OP

#define QUAT_BINARY_OP(op)                                                                                             \
    Quat operator op(T t) const                                                                                        \
    {                                                                                                                  \
        return Quat(s op t, v op t);                                                                                   \
    }                                                                                                                  \
    const Quat &operator op##=(T t)                                                                                    \
    {                                                                                                                  \
        s op## = t;                                                                                                    \
        v op## = t;                                                                                                    \
        return *this;                                                                                                  \
    }

    QUAT_BINARY_OP(*)
    QUAT_BINARY_OP(/)
#undef QUAT_BINARY_OP

    /**
     * @brief 四元数乘法
     *
     * @param[in] q
     * @return Quat&
     */
    const Quat &operator*=(const Quat &q)
    {
        *this = *this * q;
        return *this;
    }

    /**
     * @brief 四元数除法
     *
     * @param[in] q
     * @return Quat&
     */
    const Quat &operator/=(const Quat &q)
    {
        *this = (*this) / q;
        return *this;
    }

    /**
     * @brief 取反
     *
     * @return Quat
     */
    Quat operator-() const
    {
        return Quat(-s, -v);
    }

    /**
     * @brief 是否相等
     *
     * @param[in] q
     * @return true
     * @return false
     */
    bool operator==(const Quat &q) const
    {
        return (q - *this).modulus() < std::numeric_limits<T>::epsilon();
    }

    /**
     * @brief 规范化
     *
     * @return Quat&
     */
    Quat &normalize()
    {
        T len = modulus();
        ASSERT_INFO(len > FLT_MIN, "Quaternion length is zero.");
        *this /= len;
        return *this;
    }

    /**
     * @brief 规范化
     *
     * @return Quat
     */
    Quat normalized() const
    {
        T len = modulus();
        ASSERT_INFO(len > FLT_MIN, "Quaternion length is zero.");
        return *this / len;
    }

    /**
     * @brief 取模
     *
     * @return T
     */
    T modulus() const
    {
        return std::sqrt(square_modulus());
    }

    /**
     * @brief 取模平方
     *
     * @return T
     */
    T square_modulus() const
    {
        return s * s + v.square_modulus();
    }

    /**
     * @brief 序列化
     *
     * @return utils::Json
     */
    utils::Json dump_json() const
    {
        utils::Json json;
        json.append(s);
        for (std::size_t i = 0; i < 3; ++i)
            json.append(v[i]);
        return json;
    }

    /**
     * @brief 输出四元数
     *
     * @param[in] os
     * @param[in] q
     * @return std::ostream&
     */
    friend std::ostream &operator<<(std::ostream &os, const Quat &q)
    {
        os << "{" << q.s << " + " << q.v << "}";
        return os;
    }

    /**
     * @brief 点积
     *
     * @param[in] q1
     * @param[in] q2
     * @return T
     */
    friend T dot(const Quat &q1, const Quat &q2)
    {
        return q1.s * q2.s + dot(q1.v, q2.v);
    }

    /**
     * @brief 转为矩阵
     *
     * @param[in] q
     * @return Mat4x4<T>
     */
    static Mat4x4<T> to_mat4x4(const Quat &q)
    {
        auto m = Mat4x4<T>::identity();
        T s = q.s, x = q.v.x(), y = q.v.y(), z = q.v.z();
        m(0, 0) = 2 * (s * s + x * x) - 1;
        m(0, 1) = 2 * (x * y - s * z);
        m(0, 2) = 2 * (x * z + s * y);
        m(1, 0) = 2 * (x * y + s * z);
        m(1, 1) = 2 * (s * s + y * y) - 1;
        m(1, 2) = 2 * (y * z - s * x);
        m(2, 0) = 2 * (x * z - s * y);
        m(2, 1) = 2 * (y * z + s * x);
        m(2, 2) = 2 * (s * s + z * z) - 1;
        return m;
    }

    /**
     * @brief 从矩阵构造四元数
     *
     * @param[in] m
     * @return Quat
     */
    static Quat from_mat4x4(const Mat4x4<T> &m)
    {
        T trace = m(0, 0) + m(1, 1) + m(2, 2);
        if (trace > 0)
        {
            T s = std::sqrt(trace + 1) * 2;
            T inv = 1 / s;
            return Quat(0.25 * s, Vec(m(2, 1) - m(1, 2), m(0, 2) - m(2, 0), m(1, 0) - m(0, 1)) * inv);
        }
        else
        {
            // 如果 trace < 0，则 s = sqrt(1 + trace) 可能接近零，导致 0.5 / s 溢出，因此使用其它形式
            // 只需要找到对角线上最大的元素，作为正元素即可
            int i = m(0, 0) > m(1, 1) && m(0, 0) > m(2, 2) ? 0 : (m(1, 1) > m(2, 2) ? 1 : 2);
            int j = (i + 1) % 3;
            int k = (i + 2) % 3;

            Quat q;
            q.v[i] = std::sqrt(1 + m(i, i) - m(j, j) - m(k, k)) * 2;
            T inv = 1 / q.v[i];
            q.s = (m(k, j) - m(j, k)) * inv;
            q.v[j] = (m(j, i) + m(i, j)) * inv;
            q.v[k] = (m(k, i) + m(i, k)) * inv;
            q.v[i] *= 0.25;
            return q;
        }
    }

    /**
     * @brief 获得绕 axis 轴旋转 angle 角度的四元数 q
     * p -> p' = q * p * q^-1
     *
     * @param[in] axis
     * @param[in] angle
     * @return Quat
     */
    static Quat rotate(const Vec &axis, T angle)
    {
        T c = std::cos(angle / 2);
        T s = std::sin(angle / 2);
        return Quat(c, axis.normalized() * s);
    }

    /**
     * @brief 获得从 from 到 to 的旋转四元数
     *
     * @param[in] from
     * @param[in] to
     * @return Quat
     */
    static Quat rotate(const Vec &from, const Vec &to)
    {
        Vec axis = cross(from, to);
        T angle = std::acos(dot(from, to) / (from.modulus() * to.modulus()));
        return Quat::rotate(axis, angle);
    }

    /**
     * @brief 球面插值
     *
     * @param[in] q1
     * @param[in] q2
     * @param[in] t
     * @return std::function<Quat(double)>
     */
    static std::function<Quat(double)> Slerp(const Quat &q1, const Quat &q2, T t)
    {
        // angle(q1, q2) = theta，如果存在 q2 = q3 * q1，则 q3 对应的旋转角为 2 * theta
        // 因此只需要找到 qt，使得 angle(q1, qt) = (1 - t) * theta, angle(qt, q2) = t * theta
        // 于是 qt 比 q1 多旋转 2 * (1 - t) * theta，qt 比 q2 少旋转 t * theta
        T cosTheta = dot(q1, q2);

        // 如果夹角为钝角，由于 q, -q 表示同一个旋转，可以反转 q2 得到锐角
        if (cosTheta < 0)
        {
            cosTheta = -cosTheta;
            q2 = -q2;
        }

        T theta = std::acos(cosTheta);
        return [=](double t) -> Quat {
            // 如果夹角很小，可以使用线性插值 NLerp
            if (cosTheta > 1 - std::numeric_limits<T>::epsilon())
                return q1 + t * (q2 - q1);

            return (q1 * std::sin((1 - t) * theta) + q2 * std::sin(t * theta)) / std::sin(theta);
        };
    }
};

/**
 * @brief 四元数乘法
 *
 * @param[in] q
 * @return Quat
 */
template <typename T> Quat<T> operator*(const Quat<T> &lhs, const Quat<T> &rhs)
{
    return Quat<T>(lhs.s * rhs.s - dot(lhs.v, rhs.v), rhs.v * lhs.s + lhs.v * rhs.s + cross(lhs.v, rhs.v));
}

/**
 * @brief 四元数除法
 *
 * @param[in] q
 * @return Quat
 */
template <typename T> Quat<T> operator/(const Quat<T> &lhs, const Quat<T> &rhs)
{
    return lhs * Quat<T>(rhs.s, -rhs.v) / rhs.square_modulus();
}

} // namespace geom

} // namespace xi